import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import argparse
import os
from pathlib import Path

# 设置中文显示
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams["axes.unicode_minus"] = False  # 解决负号显示问题

def visualize_data(input_file, output_dir, x_col, y_col, chart_type='line', 
                  title=None, hue_col=None):
    """
    数据可视化
    :param input_file: 输入数据文件
    :param output_dir: 图表输出目录
    :param x_col: X轴列名
    :param y_col: Y轴列名
    :param chart_type: 图表类型
    :param title: 图表标题
    :param hue_col: 分组列名
    """
    try:
        # 读取数据
        ext = os.path.splitext(input_file)[1].lower()
        if ext == '.csv':
            df = pd.read_csv(input_file)
        else:
            df = pd.read_excel(input_file)
        
        # 检查列是否存在
        for col in [x_col, y_col] + ([hue_col] if hue_col else []):
            if col not in df.columns:
                print(f"数据中不存在列: {col}")
                return
        
        # 创建输出目录
        Path(output_dir).mkdir(parents=True, exist_ok=True)
        
        # 设置图表风格
        sns.set_style("whitegrid")
        plt.figure(figsize=(10, 6))
        
        # 根据类型绘制图表
        title = title or f"{chart_type}图: {y_col} vs {x_col}"
        plt.title(title)
        
        if chart_type == 'line':
            if hue_col:
                sns.lineplot(data=df, x=x_col, y=y_col, hue=hue_col, marker='o')
            else:
                sns.lineplot(data=df, x=x_col, y=y_col, marker='o')
        
        elif chart_type == 'bar':
            if hue_col:
                sns.barplot(data=df, x=x_col, y=y_col, hue=hue_col)
            else:
                sns.barplot(data=df, x=x_col, y=y_col)
        
        elif chart_type == 'scatter':
            if hue_col:
                sns.scatterplot(data=df, x=x_col, y=y_col, hue=hue_col, s=100)
            else:
                sns.scatterplot(data=df, x=x_col, y=y_col, s=100)
        
        elif chart_type == 'hist':
            sns.histplot(data=df, x=x_col, y=y_col, bins=20)
        
        else:
            print(f"不支持的图表类型: {chart_type}")
            return
        
        # 设置标签
        plt.xlabel(x_col)
        plt.ylabel(y_col)
        plt.tight_layout()
        
        # 保存图表
        output_file = os.path.join(output_dir, f"{chart_type}_{x_col}_vs_{y_col}.png")
        plt.savefig(output_file, dpi=300)
        print(f"图表已保存到: {output_file}")
        plt.close()
        
    except Exception as e:
        print(f"可视化失败: {str(e)}")

if __name__ == "__main__":
    # 检查依赖
    try:
        import pandas
        import matplotlib
        import seaborn
    except ImportError:
        print("请安装依赖: pip install pandas openpyxl matplotlib seaborn")
        exit(1)
    
    parser = argparse.ArgumentParser(description='数据可视化工具')
    parser.add_argument('input', help='输入数据文件（Excel或CSV）')
    parser.add_argument('output', help='图表输出目录')
    parser.add_argument('--x', required=True, help='X轴数据列名')
    parser.add_argument('--y', required=True, help='Y轴数据列名')
    parser.add_argument('--type', choices=['line', 'bar', 'scatter', 'hist'], 
                      default='line', help='图表类型')
    parser.add_argument('--title', help='图表标题')
    parser.add_argument('--hue', help='分组列名（用于多组数据对比）')
    
    args = parser.parse_args()
    
    visualize_data(
        input_file=args.input,
        output_dir=args.output,
        x_col=args.x,
        y_col=args.y,
        chart_type=args.type,
        title=args.title,
        hue_col=args.hue
    )
